feat: implement gemma3n text model in MLXLLM #346
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
* added to LLMModelFactory * added to MLXService * added to MLXChatExample * 4 model references from HF
|
Nice! Did you base this on #340, or did you start from scratch based on the Python implementation? |
|
Hey, good question. This implementation is like 3rd attempt and it is made from scratch based on the python source from mlx-lm. #340 was a great inspiration, since I am new to this, but sometimes it was misleading. Also, in my initial attempt I was using mlx-vlm language but it also wasn't a good reference. It all worked out once mlx-lm reference was ready. The key to a successful transpilation is to prompt it piece by piece and verify, also feeding this at the end a re-verifying the whole thing https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/converting-python |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Implementation of Gemma 3n model for MLXLLM, text only. Based on the reference implementation in mlx-lm:
ml-explore/mlx-lm#258
This code can actually help building the VLM version there #340
cc @DePasqualeOrg
Models
The original MLX weights from
mlx-vlmare not supported, only weights converted bymlx-lmare supported.I've made a new collection with Text Only MLX models, i.e.
bf16and4bitquantized using this new support.https://huggingface.co/collections/mlx-community/gemma-3n-text-only-lm-6861cf66ddc9a13102996308
Naive benchmarks
Apple M4 Max
mlx-community/gemma-3n-E4B-it-lm-bf16mlx-community/gemma-3n-E2B-it-lm-bf16mlx-community/gemma-3n-E4B-it-lm-4bitmlx-community/gemma-3n-E2B-it-lm-4bitiPhone 16 Pro
mlx-community/gemma-3n-E4B-it-lm-4bitmlx-community/gemma-3n-E2B-it-lm-4bitNotes
gelu_topk,logit_softcap) to improve performaneRMSNoScalecan be improved whenMLXFast.rmsNormis fixed (allows nil weights)Misc
Demos